import os
import logging
import pandas as pd
import numpy as np
from torch.utils.tensorboard import SummaryWriter
import shutil
from datetime import datetime

class DataLogger:
    def __init__(self, log_dir='logs', excel_dir='data', tensorboard_dir='runs'):
        self.log_dir = log_dir
        self.excel_dir = excel_dir
        self.tensorboard_dir = tensorboard_dir
        
        # 创建目录
        for dir_path in [self.log_dir, self.excel_dir, self.tensorboard_dir]:
            if os.path.exists(dir_path):
                shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        
        # 设置日志
        self.setup_logging()
        
        # 初始化TensorBoard
        self.writer = SummaryWriter(self.tensorboard_dir)
        
        # 初始化数据存储
        self.episode_data = []
        self.training_data = []
        
    def setup_logging(self):
        """设置日志配置"""
        log_file = os.path.join(self.log_dir, f'training_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log')
        
        # 清除现有的处理器
        for handler in logging.root.handlers[:]:
            logging.root.removeHandler(handler)
            
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()  # 保留控制台输出但减少内容
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    def log_episode(self, episode, total_reward, success, step_count, success_rate):
        """记录每轮训练数据"""
        episode_info = {
            'episode': episode,
            'total_reward': total_reward,
            'success': success,
            'step_count': step_count,
            'success_rate': success_rate,
            'timestamp': datetime.now()
        }
        self.episode_data.append(episode_info)
        
        # TensorBoard记录
        self.writer.add_scalar('Episode/Total_Reward', total_reward, episode)
        self.writer.add_scalar('Episode/Success', int(success), episode)
        self.writer.add_scalar('Episode/Step_Count', step_count, episode)
        self.writer.add_scalar('Episode/Success_Rate', success_rate, episode)
        
        # 日志记录（减少频率）
        if episode % 10 == 0 or success:
            self.logger.info(f"Episode {episode}: Reward={total_reward:.2f}, Success={success}, Steps={step_count}, Success_Rate={success_rate:.3f}")
    
    def log_training_step(self, episode, agent_id, actor_loss, critic_loss, step):
        """记录训练步骤数据"""
        training_info = {
            'episode': episode,
            'agent_id': agent_id,
            'actor_loss': actor_loss,
            'critic_loss': critic_loss,
            'training_step': step,
            'timestamp': datetime.now()
        }
        self.training_data.append(training_info)
        
        # TensorBoard记录
        self.writer.add_scalar(f'Agent_{agent_id}/Actor_Loss', actor_loss, step)
        self.writer.add_scalar(f'Agent_{agent_id}/Critic_Loss', critic_loss, step)
    
    def log_network_gradients(self, agent_id, network_type, gradients, step):
        """记录网络梯度信息"""
        for name, grad_info in gradients.items():
            self.writer.add_scalar(f'Agent_{agent_id}/{network_type}_Gradients/{name}_max', grad_info['max'], step)
            self.writer.add_scalar(f'Agent_{agent_id}/{network_type}_Gradients/{name}_min', grad_info['min'], step)
            self.writer.add_scalar(f'Agent_{agent_id}/{network_type}_Gradients/{name}_norm', grad_info['norm'], step)
    
    def save_to_excel(self):
        """保存数据到Excel文件"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 保存训练轮次数据
        if self.episode_data:
            episode_df = pd.DataFrame(self.episode_data)
            episode_file = os.path.join(self.excel_dir, f'episode_data_{timestamp}.xlsx')
            episode_df.to_excel(episode_file, index=False)
            self.logger.info(f"Episode data saved to {episode_file}")
        
        # 保存训练步骤数据
        if self.training_data:
            training_df = pd.DataFrame(self.training_data)
            training_file = os.path.join(self.excel_dir, f'training_data_{timestamp}.xlsx')
            training_df.to_excel(training_file, index=False)
            self.logger.info(f"Training data saved to {training_file}")
    
    def print_summary(self, episode, recent_episodes=50):
        """打印训练摘要（减少控制台输出）"""
        if len(self.episode_data) >= recent_episodes:
            recent_data = self.episode_data[-recent_episodes:]
            avg_reward = np.mean([d['total_reward'] for d in recent_data])
            success_rate = np.mean([d['success'] for d in recent_data])
            avg_steps = np.mean([d['step_count'] for d in recent_data])
            
            print(f"\n=== Episode {episode} Summary (Last {recent_episodes} episodes) ===")
            print(f"Average Reward: {avg_reward:.2f}")
            print(f"Success Rate: {success_rate:.3f}")
            print(f"Average Steps: {avg_steps:.1f}")
            
            self.logger.info(f"Summary - Avg Reward: {avg_reward:.2f}, Success Rate: {success_rate:.3f}, Avg Steps: {avg_steps:.1f}")
    
    def close(self):
        """关闭记录器"""
        self.save_to_excel()
        self.writer.close()
        
        # 最终统计
        if self.episode_data:
            total_episodes = len(self.episode_data)
            total_success = sum([d['success'] for d in self.episode_data])
            avg_reward = np.mean([d['total_reward'] for d in self.episode_data])
            
            print(f"\n=== Training Completed ===")
            print(f"Total Episodes: {total_episodes}")
            print(f"Overall Success Rate: {total_success/total_episodes:.3f}")
            print(f"Average Reward: {avg_reward:.2f}")
            
            self.logger.info(f"Training completed - Episodes: {total_episodes}, Success Rate: {total_success/total_episodes:.3f}, Avg Reward: {avg_reward:.2f}")